import os
import time
import copy
import threading
from typing import List, Optional
from pathlib import Path
import threading
from collections import deque
import numpy as np
import rospy
import tf2_ros
import cv2
from cv_bridge import CvBridge
import message_filters
from sensor_msgs.msg import Image, CameraInfo
COLOR_TOPICS = ["/camera_right/color/image_raw"]
DEPTH_TOPICS = ["/camera_right/aligned_depth_to_color/image_raw"]
POINTS_TOPICS = ["/camera_right/depth/points"]
COLOR_INFO_TOPICS = ["/camera_right/color/camera_info"]
DEPTH_INFO_TOPICS = ["/camera_right/aligned_depth_to_color/camera_info"]


def stamp_to_sec(stamp):
    return stamp.to_sec()


class CameraParam:
    def __init__(self, frame_id: str, intrinsic_matrix: np.ndarray, resolution: Optional[tuple] = None, distortion_coeff: Optional[np.ndarray] = None):
        self.frame_id = frame_id
        self.k = intrinsic_matrix
        self.resolution = resolution
        self.distortion_coeff = distortion_coeff

    def __str__(self):
        return (
            f"Camera Name: {self.frame_id}\n"
            f"Resolution: {self.resolution[0]}x{self.resolution[1]}\n" if resolution else ""
            f"Intrinsic Matrix:\n{self.k}\n"
            f"Distortion Coefficients:\n{self.distortion_coeff}"
        )

    def __repr__(self):
        return self.__str__()

    def get_intrinsic(self):
        return self.k

    def set_intrinsic(self, camera_intrinsic):
        self.k = camera_intrinsic


class ObservBuffer:
    def __init__(self, max_len=10):
        self.buffer = deque(maxlen=max_len)
        self.lock = threading.Lock()

    def push(self, observation):
        with self.lock:
            self.buffer.append(observation)

    def pop(self):
        with self.lock:
            if len(self.buffer) > 0:
                return self.buffer.pop()
            else:
                return None

    def len(self):
        with self.lock:
            return len(self.buffer)

    def latest(self):
        with self.lock:
            if len(self.buffer) > 0:
                return copy.deepcopy(self.buffer[-1])
            else:
                return None

    def clear(self):
        with self.lock:
            self.buffer.clear()

    def get(self, index):
        with self.lock:
            if index < len(self.buffer):
                return copy.deepcopy(self.buffer[index])
            else:
                return None

    def get_all(self):
        with self.lock:
            return copy.deepcopy(self.buffer)


class ParamManager():
    def __init__(self):
        self.params = {}
        self.params['use_point_cloud'] = rospy.get_param("~use_point_cloud", False)
        self.params['color_topics'] = rospy.get_param("~color_topics", COLOR_TOPICS)
        self.params['depth_topics'] = rospy.get_param("~depth_topics", DEPTH_TOPICS)
        self.params['color_info_topics'] = rospy.get_param("~color_info_topics", COLOR_INFO_TOPICS)
        self.params['depth_info_topics'] = rospy.get_param("~depth_info_topics", DEPTH_INFO_TOPICS)
        self.params['queue_size'] = rospy.get_param("~queue_size", 10)
        self.params['slop'] = rospy.get_param("~slop", None)
        self.params['output_dir'] = rospy.get_param("~output_dir", "./outputs")

    def get_params(self):
        return self.params

    def get_param(self, name):
        return self.params.get(name, None)

    def get_param_ros(self, name, default=None):
        self.params[name] = rospy.get_param(name, default)
        return self.params[name]

    def set_params(self, params):
        self.params.update(params)

    def set_param(self, name, value):
        self.params[name] = value


class TransformBuffer:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls, *args, **kwargs):
        ''' singleton pattern '''
        with cls._lock:
            if not cls._instance:
                cls._instance = super(TransformBuffer, cls).__new__(cls)
        return cls._instance

    def __init__(self, cache_duration=20.0):
        self.tf_buffer = tf2_ros.Buffer(cache_time=rospy.Duration(cache_duration))
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

    def query(self, target_frame, source_frame, time_stamp=rospy.Time(0), duration=0.1):
        try:
            tf = self.tf_buffer.lookup_transform(target_frame, source_frame, time_stamp, rospy.Duration(duration))
            return tf
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logerr(f"Error looking up transform: {e}")
            return None

    def query_full(self, target_frame, target_time, source_frame, source_time, fixed_frame='map', duration=0.1):
        try:
            tf = self.tf_buffer.lookup_transform_full(
                target_frame=target_frame,
                target_time=target_time,
                source_frame=source_frame,
                source_time=source_time,
                fixed_frame=fixed_frame,
                timeout=rospy.Duration(duration)
            )
            return tf
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logerr(f"Error looking up transform: {e}")
            return None


class MultiTopicSynchronizer:
    """
    Generic class that implements time synchronization for multiple ROS topic messages.
    """
    def __init__(self, callback=None):
        self.callback = callback

    def init(self, topics_and_msgtypes, queue_size=10, slop=None):
        """
        Initialize the MultiTopicSynchronizer.

        Parameters:
            topics_and_msgtypes (list[tuple[str, msg_type]]): A list of topics to be synchronized along with their corresponding message types.
            queue_size (int): The size of the synchronizer message queue.
            slop (float): The allowable timestamp deviation in seconds.If None, use ExactTimeSynchronizer instead of ApproximateTimeSynchronizer.
        """
        # Create a list of subscribers
        self.subscribers = []
        for topic, message_type in topics_and_msgtypes:
            subscriber = message_filters.Subscriber(topic, message_type)
            self.subscribers.append(subscriber)
        rospy.loginfo("The slop for TimeSynchronizer is {}".format(slop))
        if slop is None:
            self.synchronizer = message_filters.TimeSynchronizer(self.subscribers, queue_size)
        else:
            self.synchronizer = message_filters.ApproximateTimeSynchronizer(self.subscribers, queue_size, slop)
        return self

    def start(self):
        """
        Start the MultiTopicSynchronizer.
        """
        self.synchronizer.registerCallback(self.sync_callback)

    def sync_callback(self, *msgs):
        """
        Synchronized message callback function.

        Parameters:
            *msgs: A list of synchronized messages in the order of the topics list.
        """
        # Process the synchronized messages here.
        if self.callback is not None:
            self.callback(*msgs)


class CameraInfoSubscriber:
    def __init__(self, topics):
        self.num_camerainfo = len(topics)
        self.camera_infos = {}
        if isinstance(topics, dict):
            self.camera_info_topics = list(topics.values())
            self.camera_info_names = list(topics.keys())
        elif isinstance(topics, list):
            self.camera_info_topics = topics
            self.camera_info_names = [i for i in range(self.num_camerainfo)]

        self.current_info_idx = 0
        self.current_info_subscription = None
        self.lock = threading.Lock()
        self.subscribe_to_next_info()

    def subscribe_to_next_info(self):
        if self.current_info_idx < self.num_camerainfo:
            self.current_info_subscription = rospy.Subscriber(
                self.camera_info_topics[self.current_info_idx],
                CameraInfo,
                self.camera_info_callback,
                queue_size=1
            )
            rospy.loginfo(f'Waiting for camera {self.current_info_idx} info...')
        else:
            rospy.loginfo('All camera info received')

    def camera_info_callback(self, msg):
        rospy.loginfo("color info received")
        with self.lock:
            self.camera_infos[self.camera_info_names[self.current_info_idx]] = CameraParam(str(self.current_info_idx), np.array(msg.K).reshape(3, 3))
            rospy.loginfo("Intrinsic matrix: \n{}".format(self.camera_infos[self.camera_info_names[self.current_info_idx]].get_intrinsic()))
        # Destroy current subscription
        self.current_info_subscription.unregister()
        self.current_info_subscription = None
        # Subscribe to the next camera topic
        self.current_info_idx += 1
        self.subscribe_to_next_info()

    def len(self):
        with self.lock:
            return len(self.camera_infos)

    def get_camera_info(self, camera_name):
        with self.lock:
            if camera_name in self.camera_infos:
                return self.camera_infos[camera_name]
        return None

    def get_camera_infos(self):
        with self.lock:
            return copy.deepcopy(self.camera_infos)


class ColorDepthSyncer(MultiTopicSynchronizer):
    def __init__(self, parameters, process_observation=None):
        if process_observation is None:
            self.process_observation = lambda color_msgs, depth_msgs: None
        else:
            self.process_observation = process_observation

        def sync_callback(*msgs):
            rospy.loginfo("synchronized messages received")
            color_msgs = msgs[:self.num_color]
            depth_msgs = msgs[self.num_color:]
            self.process_observation(color_msgs, depth_msgs)

        super().__init__(callback=sync_callback)
        self.use_point_cloud = parameters.get('use_point_cloud', False)
        color_topics = parameters.get('color_topics', COLOR_TOPICS)
        depth_topics = parameters.get('depth_topics', DEPTH_TOPICS)
        queue_size = parameters.get('queue_size', 10)
        slop = parameters.get('slop', None)
        self.num_color = len(color_topics)
        if self.use_point_cloud:
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, PointCloud2) for topic in depth_topics]
        else:
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, Image) for topic in depth_topics]

        rospy.loginfo(f"subscribed topics and msgtypes: {topics_and_msgtypes}")
        super().init(topics_and_msgtypes, queue_size=queue_size, slop=slop)
        super().start()


class ColorDepthSaver(ColorDepthSyncer):
    def __init__(self, parameters: dict):
        def process_observation(color_msgs, depth_msgs):
            rospy.loginfo(f"processing observation")
            start_time = time.time()
            time_stamp, self.observations = self.save_observation(color_msgs, depth_msgs)
            end_time = time.time()
            print(f"convert observation consumes time: {end_time - start_time:.3f}s")
            time_str = self.timestamp_to_string(time_stamp)  # 使用新的方法
            rospy.loginfo(f"Timestamp as string: {time_str}")
        super().__init__(parameters, process_observation=process_observation)
        self.cv_bridge = CvBridge()
        self.output_dir = parameters.get("output_dir", "./outputs")
        os.makedirs(self.output_dir, exist_ok=True)
        self.color_dir = os.path.join(self.output_dir, "rgb")
        os.makedirs(self.color_dir, exist_ok=True)
        self.depth_dir = os.path.join(self.output_dir, "depth")
        os.makedirs(self.depth_dir, exist_ok=True)

    def __del__(self):
        for intrinsic in self.camera_intrinsics:
            file_name = os.path.join(self.output_dir, "cam_K.txt")
            np.savetxt(file_name, intrinsic, fmt="%.18f")

    def save_observation(self, color_msgs, depth_msgs):
        color_images = []
        depth_images = []
        for i in range(len(color_msgs)):
            cv_color_image = self.cv_bridge.imgmsg_to_cv2(color_msgs[i], desired_encoding="bgr8")
            timestamp = self.timestamp_to_string(color_msgs[i].header.stamp)
            cv2.imwrite(f"{self.color_dir}/{timestamp}.png", cv_color_image)
            color_images.append(cv_color_image)

        for i in range(len(depth_msgs)):
            cv_depth_image = self.cv_bridge.imgmsg_to_cv2(depth_msgs[i], desired_encoding="16UC1")
            timestamp = self.timestamp_to_string(depth_msgs[i].header.stamp)
            cv2.imwrite(f"{self.depth_dir}/{timestamp}.png", cv_depth_image)
            depth_images.append(cv_depth_image)

        return color_msgs[0].header.stamp, list(zip(color_images, depth_images))

    def timestamp_to_string(self, timestamp):
        """
        Convert rospy.Time to a formatted string in seconds.

        Parameters:
            timestamp (rospy.Time): The timestamp to convert.

        Returns:
            str: The formatted timestamp string in seconds.
        """
        # Convert rospy.Time to seconds
        timestamp_sec = timestamp.to_sec()
        # Format seconds to string
        return f"{timestamp_sec:.6f}"  # 保留6位小数


class ColorDepthBuffer(ObservBuffer):
    def __init__(self, parameters: dict):
        super().__init__(max_len=10)
        def process_observation(color_msgs, depth_msgs):
            rospy.loginfo(f"processing observation")
            # 提取所有时间戳
            timestamps = [stamp_to_sec(msg.header.stamp) for msg in color_msgs + depth_msgs]
            # 计算时间戳之间的差异
            max_diff = max(timestamps) - min(timestamps)
            # time_stamp, self.observations = self.save_observation(color_msgs, depth_msgs)
            rospy.loginfo(f"max difference between Timestamp: {max_diff:.6f}")
            self.push((color_msgs, depth_msgs))
            rospy.loginfo(f"save observation {self.len()}")
        self.color_depth_syncer = ColorDepthSyncer(parameters, process_observation=process_observation)


def main():
    node_name = 'color_depth_saver'
    rospy.init_node(node_name, anonymous=False)
    rospy.loginfo("Starting ColorDepthSyncer node...")
    param_manager = ParamManager()
    parameters = param_manager.get_params()
    color_depth_saver = ColorDepthBuffer(parameters)
    camera_info_sub = CameraInfoSubscriber(COLOR_INFO_TOPICS)
    rospy.spin()


if __name__ == "__main__":
    main()
